from captcha.image import ImageCaptcha
import random
from PIL import Image
import numpy as np
import tensorflow as tf

number = list( map(lambda x:str(x) ,range(10)))
alphabet = list( map(lambda x:chr(x) , range(ord("a"),ord("z")+1)) )
ALPHABET = list( map(lambda x:chr(x) , range(ord("A"),ord("Z")+1)) )

CHAR_SET = number+alphabet+ALPHABET
SAVE_PATH = r"D:\pycharmProjects\tf2.0cnn验证码识别\keras_cnn"
CHAR_SET_LEN = len(CHAR_SET)
IMAGE_HIGH = 60
IMAGE_WIDTH = 160
num_epochs =20
learning_rate = 0.001
batch_size = 512

def random_captcha_text(char_set=None,captcha_size = 4):
    if char_set is None :
        char_set = CHAR_SET
    captcha_text=[]
    for i in range(captcha_size):
        c = random.choice(char_set)
        captcha_text.append(c)
    return captcha_text

def gen_char_text_and_image(width=160,height=60,char_set = CHAR_SET):
    image = ImageCaptcha(width=width,height=height)
    captcha_text = random_captcha_text(char_set)
    captcha_text = ''.join(captcha_text)
    captcha = image.generate(captcha_text)

    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)
    return captcha_text , captcha_image

text,image = gen_char_text_and_image(char_set = CHAR_SET)
MAX_CAPTCHA = len(text)
print('char_ser_len=%s,max_captcha=%s'%(CHAR_SET_LEN,MAX_CAPTCHA))

def image2gray(img):
    if len(img.shape) >2 :
        gray = np.mean(img,-1)
        return gray
    else:
        return img

def text2vec(text):
    vector = np.zeros([MAX_CAPTCHA,CHAR_SET_LEN])
    for i,c in enumerate(text):
        idx = CHAR_SET.index(c)
        vector[i][idx]=1.0
    return vector

def vec2text(vector):
    text = []
    # for i,c in enumerate(vector):
        # l = c.astype(np.bool)
        # l_ = np.array(CHAR_SET)[l][0]
        # text.append(l_)
    for i, c in enumerate(vector):
        text.append(CHAR_SET[c])
    return "".join(text)

def get_next_batch(batch_size = 128):
    batch_x = np.zeros([batch_size,IMAGE_HIGH,IMAGE_WIDTH,1])
    batch_y = np.zeros([batch_size,MAX_CAPTCHA,CHAR_SET_LEN])

    def wrap_gen_text_and_img():
        while True :
            text,img = gen_char_text_and_image(char_set=CHAR_SET)
            if img.shape==(60,160,3):
                return text,img

    for i in range(batch_size):
        text,img = wrap_gen_text_and_img()
        img = tf.reshape(image2gray(img),[IMAGE_HIGH,IMAGE_WIDTH,1])
        batch_x[i] = img
        batch_y[i] = text2vec(text)

    return batch_x,batch_y

def crack_captcha_cnn():
    model = tf.keras.Sequential()

    # 160*60*1  --> 80*30*32
    model.add(tf.keras.layers.Conv2D(32,(3,3)))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2,2),strides=2))

    # 80*30*32  --> 40*15*64
    model.add(tf.keras.layers.Conv2D(64, (5, 5)))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2, 2), strides=2))

    # 40*15*64  --> 20*8*128
    model.add(tf.keras.layers.Conv2D(128, (5, 5)))
    model.add(tf.keras.layers.PReLU())
    model.add(tf.keras.layers.MaxPool2D((2, 2), strides=2))

    model.add(tf.keras.layers.Dropout(0.3))

    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(MAX_CAPTCHA*CHAR_SET_LEN))
    model.add(tf.keras.layers.Reshape([MAX_CAPTCHA, CHAR_SET_LEN]))

    model.add(tf.keras.layers.Softmax())

    return model

def train():
    try:
        model = tf.keras.models.load_model(SAVE_PATH+'\model')
    except Exception as e:
        print("############Exception",e)
        model = crack_captcha_cnn()

    model.compile(optimizer = tf.keras.optimizers.Adam(0.00001),metrics=['accuracy'],loss='categorical_crossentropy')

    for times in range(5000):
        batch_x,batch_y = get_next_batch(512)
        print('批次数=', times, ' batch_x.shape=', batch_x.shape, ' batch_y.shape=', batch_y.shape)

        model.fit(batch_x, batch_y, epochs=4)
        print("y预测=\n",np.argmax(model.predict(batch_x),axis=2))
        print("y实际=\n", np.argmax(batch_y, axis=2))

        if 0 == times%10 and times != 0:
            print("在第%s批保存模型"%(times))
            model.save(SAVE_PATH+'\model')

def predict():
    model = tf.keras.models.load_model(SAVE_PATH+'\model')
    success = 0
    count = 100
    for _ in range(count):
        data_x,data_y = get_next_batch(1)
        predict_value = model.predict(data_x)

        data_y = vec2text(np.argmax(data_y,axis=2)[0])
        prediction_value = vec2text(np.argmax(predict_value, axis=2)[0])

        if data_y.upper() == prediction_value.upper():
            print("y预测=", prediction_value, "y实际=", data_y, "预测成功。")
            success += 1
        else:
            print("y预测=", prediction_value, "y实际=", data_y, "预测失败。")
        print("===预测", count, "次,", "成功率=", success / count)


if __name__ == '__main__':
    # train()
    predict()